# 作者：hehung
# 联系：1398660197@qq.com
# 微信：hehung95

import socket
import sys
import cv2
import numpy as np
import sys
import mediapipe as mp
import math
import struct


SERVER_IP_PORT_RECV = ('192.168.6.101', 6789)
SERVER_IP_PORT_SEND = ('192.168.6.105', 6790)
CHUNK_SIZE = 8192

wrist_x_angle = 0.0
wrist_y_angle = 0.0
hand_dist = 0.0
hand_x_norm = 0.0
hand_y_norm = 0.0
catch = 0.0
steer = [7.5, 7.5, 7.5, 7.5, 7.5, 7.5]
steer_filter = [[7.5] * 10 for _ in range(6)]

FILTER_NUM = 10

class ImgClass:
    def __init__(self):
        # 获取摄像机
        self.cap = cv2.VideoCapture(0)

    def read_cap(self):
        # 读取摄像头采集的视频
        ret, frame = self.cap.read()
        return frame

    # 显示图片
    def show_img(self, img):
        cv2.imshow('Receiver', img)
        c = cv2.waitKey(1)
        # 如果按下ESC则退出
        if c == 27:
            self.close_img()
            return 0
        else:
            return 1

    # 关闭显示
    def close_img(self):
        self.cap.release()
        cv2.destroyAllWindows()

    # 获取图片大小
    def get_img_size(self, img):
        height, width = img.shape[:2]

        return [height, width]

    # 将图片转换字节
    def img_2_bytes(self, img):
        img_encode = cv2.imencode('.jpg', img)[1]
        data_encode = np.array(img_encode)
        data = data_encode.tobytes()
        return data
    
    # 将字节转换为图片
    def bytes_2_img(self, bytes):
        nparr = np.frombuffer(bytes, np.uint8)
        img_decode = cv2.imdecode(nparr, flags=cv2.IMREAD_COLOR)
        return img_decode

class SocketClass:
    def __init__(self, ip_port):
        try:
            # 创建一个基于 IPv4 和 UDP 协议的套接字对象
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 650000)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 650000)
            # 设置套接字选项，允许地址重用
            # self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_ip_port = ip_port
            print (self.server_socket.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF))
        except socket.error as msg:
            print(msg)
            print(sys.exit(1))
        
    # 绑定IP和端口
    def socket_bind(self):
        # 将套接字绑定到指定的主机和端口
        self.server_socket.bind(self.server_ip_port)

    # 分包发送图片
    def send_in_chunks(self, img_bytes):
        global steer
        for i in range(6):
            # 取小数点位数为2
            steer[i] = round(steer[i], 2)
            # print (steer[i])
        data = struct.pack("!6f", steer[0], steer[1], steer[2], steer[3], steer[4], steer[5]) + img_bytes
        self.server_socket.sendto(data, self.server_ip_port)

    # 分包接收图片
    def recv_in_chunks(self):
        img_data, addr = self.server_socket.recvfrom(600000)
        return img_data

# 手势识别
mp_hands = mp.solutions.hands
hands = mp_hands.Hands()
mp_drawing = mp.solutions.drawing_utils

def average_filter(steer_num, value):
    for i in range(FILTER_NUM):
        if (i < FILTER_NUM - 1):
            steer_filter[steer_num][i] = steer_filter[steer_num][i+1]
        else:
            steer_filter[steer_num][i] = value
        steer[steer_num] = steer[steer_num] + steer_filter[steer_num][i]
    steer[steer_num] = steer[steer_num] / FILTER_NUM

# 计算手指到手腕的距离
def calc_finger_dist(landmarks, finger_tip_index1, finger_tip_index2):
    x1, y1 = landmarks.landmark[finger_tip_index1].x, landmarks.landmark[finger_tip_index1].y
    x2, y2 = landmarks.landmark[finger_tip_index2].x, landmarks.landmark[finger_tip_index2].y
    
    dist = math.hypot(x2 - x1, y2 - y1)
    return dist


# 手势处理
def mediapipe_process(frame):
    global wrist_x_angle
    global wrist_y_angle
    global hand_dist
    global hand_x_norm
    global hand_y_norm
    global catch
    # 将BGR转换为RGB
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 手势识别
    results = hands.process(image)
    # 再将RGB转回BGR
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    if results.multi_hand_landmarks:
        hand_landmarks = results.multi_hand_landmarks[0]
        if True:
        # for hand_landmarks in results.multi_hand_landmarks:
            ########################################################
            # # 计算手势切斜角度
            # wrist_x, wrist_y = hand_landmarks.landmark[mp_hands.HandLandmark.WRIST].x, hand_landmarks.landmark[mp_hands.HandLandmark.WRIST].y
            # index_x, index_y = hand_landmarks.landmark[mp_hands.HandLandmark.INDEX_FINGER_TIP].x, hand_landmarks.landmark[mp_hands.HandLandmark.INDEX_FINGER_TIP].y  
            # # 计算掌心倾斜角度
            # v = [index_x - wrist_x, index_y - wrist_y]  
            # wrist_y_angle = math.atan2(v[1], v[0]) * 180 / math.pi
            # # print(f"Tilt Angle: {wrist_y_angle}")

            ########################################################
            # 计算手掌到摄像头的距离
            # dist = np.round(hand_landmarks.landmark[0].z * 1000000000, 5)
            v1 = hand_landmarks.landmark[0] # 手腕点
            v2 = hand_landmarks.landmark[1] # 大拇指对下面的点
            hand_dist = math.fabs(math.sqrt((v1.x - v2.x)**2 + (v1.y - v2.y)**2))
            if (hand_dist <= 0.04):
                hand_dist = 0.04
            elif (hand_dist >= 0.08):
                hand_dist = 0.08
            hand_dist = (hand_dist - 0.08) / (0.04 - 0.08)
            if (hand_dist < 0):
                hand_dist = 0
            elif (hand_dist > 1):
                hand_dist = 1
            # print (f'distance: {hand_dist}')
            
            ########################################################
            # 手在摄像头中的位置
            # 假设检测结果如下
            wrist = hand_landmarks.landmark[0] 
            # 归一化坐标
            hand_x_norm = np.round(wrist.x, 5)
            if (hand_x_norm > 10):
                hand_x_norm = 1
            elif (hand_x_norm < 0):
                hand_x_norm = 0
            hand_y_norm = np.round(wrist.y, 5)
            if (hand_y_norm < 0.4):
                hand_y_norm = 0.4
            elif (hand_y_norm > 1):
                hand_y_norm = 1
            hand_y_norm = (hand_y_norm - 0.4) * 1.6
            
            # print ('Hand position:' + str(hand_x_norm) + '     ' + str(hand_y_norm))

            ########################################################
            # 计算每个手指到大拇指的距离
            # 0 - 手腕
            # 4 - 拇指指尖
            # 12 - 中指指尖
            # 16 - 小拇指指尖
            # thumb_dist = np.round(calc_finger_dist(hand_landmarks, 0, 4), 5)
            middle_dist = np.round(calc_finger_dist(hand_landmarks, 0, 12), 5)
            # ring_dist = np.round(calc_finger_dist(hand_landmarks, 0, 16), 5)
            # 补偿 0.5 ~ 0.8，标定值(手掌伸开)
            # 补偿 0.28 ~ 0.12 标定值(手掌握拳)
            # 归一化中指指尖到手腕距离，用于判断是否握拳
            # if (middle_dist < 0.2):
            #     middle_dist = 0.2
            # elif (middle_dist > 0.8):
            #     middle_dist = 0.8
            # 理论距离公式（手掌伸开时）
            middle_dist_norm_extend = (-0.24 * hand_dist + 0.44)
            # 理论距离公式（手掌握拳时）
            middle_dist_norm_fist   = (-0.08 * hand_dist + 0.2)
            # 差值
            if (middle_dist > middle_dist_norm_extend):
                middle_dist = middle_dist_norm_extend
            elif (middle_dist < middle_dist_norm_fist):
                middle_dist = middle_dist_norm_fist
            # 归一化 0 ~ 1
            fist_v_per = (middle_dist - middle_dist_norm_fist) / (middle_dist_norm_extend - middle_dist_norm_fist)
            catch = 1 - fist_v_per

            # print ("Finger distances:", middle_dist, hand_dist)

            ########################################################
            # 计算手掌绕手腕翻转角度
            # 拇指指关节点  
            thumb_mcp_x, thumb_mcp_y = hand_landmarks.landmark[3].x * 300, hand_landmarks.landmark[3].y * 400
            # 小指指根节点
            pinky_mcp_x, pinky_mcp_y = hand_landmarks.landmark[17].x * 300, hand_landmarks.landmark[17].y * 400    
            # 计算角度
            dx = pinky_mcp_x - thumb_mcp_x  
            dy = pinky_mcp_y - thumb_mcp_y
            wrist_x_angle = math.atan2(dy, dx) * 180 / math.pi
            if (wrist_x_angle < -80):
                wrist_x_angle = -80
            elif (wrist_x_angle > 80):
                wrist_x_angle = 80
            wrist_x_angle = wrist_x_angle + 80
            wrist_x_angle = 160 - wrist_x_angle
            # range 0 ~ 160
            # print(f"Palm rotation angle: {wrist_x_angle} degrees")

            mp_drawing.draw_landmarks(
                image,
                hand_landmarks,
                mp_hands.HAND_CONNECTIONS) # 用于指定地标如何在图中连接。
    return image

def pose_calc():
    global wrist_x_angle
    global wrist_y_angle
    global hand_dist
    global hand_x_norm
    global hand_y_norm
    global catch
    global steer

    # 控制逻辑为0~100%占空比，100占空比对应的20ms
    # SG90舵机PWM控制： 0度周期0.5ms，对应占空比为2.5%
    #                  180度周期2.5ms，对应占空比12.5%
    # 旋转计算
    # 控制舵机旋转，中间为90度，对应SG90舵机PWM控制占空比7.5%
    # 舵机0度到180度，对应2.5%~12.5%
    hand_x_norm = 1 - hand_x_norm
    steer[0] = 2.5 + (hand_x_norm) * 10
    # 机械臂夹子夹的力度计算
    # 取值为0 ~ 0.2，0对应没握拳，0.2对应全握
    # 计算公式如下
    steer[5] = 7.5 + (catch) * 5
    # 翻转, 范围：0~100
    steer[4] = 2.5 + (wrist_x_angle) / 18
    # 上下
    # 取值范围 0~1
    hand_dist = 1 - hand_dist
    steer[1] = 2.5 + (hand_dist) * 10

    # 前后
    steer[2] = 2.5 + (hand_y_norm) * 10

    # average_filter(0, steer[0])
    # average_filter(1, steer[1])
    # average_filter(2, steer[2])
    # average_filter(4, steer[4])
    # average_filter(5, steer[5])

    # for i in range(6):
    #     # 取小数点位数为2
    #     steer[i] = round(steer[i], 2)
    #     print (steer[i], end = ' -- ')
    # print ()


if __name__ == '__main__':
    print ('Start Server')
    img_handle = ImgClass()
    udp_socket_recv = SocketClass(SERVER_IP_PORT_RECV)
    udp_socket_recv.socket_bind()
    udp_socket_send = SocketClass(SERVER_IP_PORT_SEND)
    print ('Server started')
    while True:
        img_data = udp_socket_recv.recv_in_chunks()
        # 转换为图片数组
        img = img_handle.bytes_2_img(img_data)
        # 手势处理
        img = mediapipe_process(img)
        pose_calc()
        # 显示图片
        if (0 == img_handle.show_img(img)):
            break
        # 将图片数据转换为字节方便发送给服务器
        # img = cv2.resize(img, (400, 300)) 
        frame_data = img_handle.img_2_bytes(img)
        # 将图片发送回客户端
        udp_socket_send.send_in_chunks(frame_data)
    print ('Exit by manual')
